# %%
import cv2
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.decomposition import PCA
from scipy.interpolate import make_interp_spline

def GetData(index):
    '''
    Read grayscale images in the data set
    '''
    fp = "E:\Data\orl_faces"
    X, y = pd.DataFrame(np.array([])), []

    for i in range(1, 41):
        for j in index:
            imagPath = fp + '\s' + str(i) + '\\' + str(j) + '.pgm'
            newimag = cv2.imread(imagPath, cv2.IMREAD_GRAYSCALE)
            # Convert the image to 1 dimensional
            newimag = newimag.reshape(1, newimag.size)
            X = pd.concat([X, pd.DataFrame(newimag)], axis=0)
            y.append(i) # label
    y = np.array(y).ravel()
    return X, y

# %%
def StdData(X_train, X_test):
    '''
    Standardize data
    '''
    stdScaler = StandardScaler().fit(X_train)
    X_train = stdScaler.transform(X_train)
    X_test = stdScaler.transform(X_test)
    return X_train, X_test

# %%
def PcaDimRedu(X_train, X_test, comNum=20):
    '''
    Use PCA to reduce dimensionality of data
    '''
    pca = PCA(n_components=comNum).fit(X_train) 
    X_train = pca.transform(X_train) 
    X_test = pca.transform(X_test)
    return X_train, X_test

# %%
def GetCorrRate(y_test, y_test_pred):
    '''
    Calculate the correct rate
    '''
    corrNum = 0
    for i in range(len(y_test)):
        if y_test[i] == y_test_pred[i]:
            corrNum = corrNum + 1
    return corrNum / len(y_test)

# %%
def plotComponent_correctRate_runningTime(comNum, corrRates, runningTime):
    '''
    Draw a graph with Component Number on the horizontal axis 
    and Correct Rate on the vertical axis
    '''
    x = np.array(comNum)
    y1 = np.array(corrRates)
    y2 = np.array(runningTime)
    x_new = np.linspace(x.min(), x.max(), 500)
    y1 = make_interp_spline(x, y1)(x_new)
    y2 = make_interp_spline(x, y2)(x_new)
    
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.plot(x_new, y1, 'r', label = 'correct rate')
    ax1.legend(loc=1)
    ax1.set_ylabel("Correct Rate")

    ax2 = ax1.twinx()
    ax2.plot(x_new, y2, 'g', label = "Running Time")
    ax2.legend(loc=2)
    ax2.set_ylabel("Running Time(s)")

    ax1.set_xlabel("The number of component")
    return None

# %%
def plotComponent_correctRate(comNum, allCorrRates, kernelNames):
    '''
    Draw an image of the change in correct rate when the kernel function 
    and the number of principal components are different
    '''
    plt.xlabel("The number of commponet")
    plt.ylabel("Correct Rate")
    x = np.array(comNum)
    x_new = np.linspace(x.min(), x.max(), 500)
    for i in range(len(kernelNames)):
        y = np.array(allCorrRates[i])
        y_smooth = make_interp_spline(x, y)(x_new)
        plt.plot(x_new, y_smooth, label = kernelNames[i])
    plt.legend(kernelNames)
    return None

# %%
def componentImpact(X_train, y_train, X_test, y_test):
    '''
    Explore the impact of the kernel function on the results
    '''
    comNum = [10, 15, 20, 50, 100, 150, 200] # The numbers of components
    corrRates, runningTime = [], []

    for i in comNum:
        startTime = time.time()
        # PCA dimensionality reduction(Specify the number of principal component)
        X_train_std, X_test_std = PcaDimRedu(X_train, X_test, i)
        # Training SVM model(polynomial kernel function)
        svm = SVC(kernel='poly').fit(X_train_std, y_train)
        # Use the model to classify the test set
        y_test_pred = svm.predict(X_test_std)
        # Calculate the correct rate of classification
        corrRate = GetCorrRate(y_test, y_test_pred)
        endTime = time.time()
        corrRates.append(corrRate)
        runningTime.append(endTime - startTime)
    plotComponent_correctRate_runningTime(comNum, corrRates, runningTime)
    return None

# %%
def kernelImpact(X_train, y_train, X_test, y_test):
    '''
    Explore the influence of the number of principal components on the results
    '''
    kernelNames = ['poly', 'linear', 'rbf', 'sigmoid']
    comNum = [10, 15, 20, 50, 100, 150, 200] # The numbers of components
    allCorrRate = []
    for i in range(len(kernelNames)):
        corrRates = []
        for j in comNum:
            X_train_std, X_test_std = PcaDimRedu(X_train, X_test, j)
            svm = SVC(kernel=kernelNames[i]).fit(X_train_std, y_train)
            y_test_pred = svm.predict(X_test_std)
            corrRate = GetCorrRate(y_test, y_test_pred)
            corrRates.append(corrRate)
        allCorrRate.append(corrRates)
    plotComponent_correctRate(comNum, allCorrRate, kernelNames)
    return None

# %%
if __name__ == '__main__':
    # Get Data
    X_train, y_train = GetData(range(1, 6))
    X_test, y_test = GetData(range(6, 11))
    # Standardize data
    X_train, X_test = StdData(X_train, X_test)
    
    # Explore the influence of the number of principal components on the results
    # componentImpact(X_train, y_train, X_test, y_test)
    
    # Explore the impact of the kernel function on the results
    kernelImpact(X_train, y_train, X_test, y_test)

    # Exchange test set and training set for testing
    # kernelImpact(X_test, y_test, X_train, y_train)


